import jieba
import pandas as pd
import jieba.posseg as psg
from imageio import imread
from wordcloud import get_single_color_func
import matplotlib.pyplot as plt
import wordcloud
raw = pd.read_csv("金庸-射雕英雄传txt精校版.txt",names=['txt'], sep='aaa', encoding="GBK" ,engine='python')
def m_head(tem_str):
    return tem_str[:1]
def m_mid(tmp_str):
    return tmp_str.find("回 ")
raw['head'] = raw.txt.apply(m_head)
raw['mid'] = raw.txt.apply(m_mid)
raw['len'] = raw.txt.apply(len)
raw.head(50)
chapnum = 0
for i in range(len(raw)):
    if raw['head'][i] == "第" and raw['mid'][i] > 0 and raw['len'][i] < 30:
        chapnum += 1
    if chapnum >= 40 and raw['txt'][i] == "附录一：成吉思汗家族":
        chapnum = 0
    raw.loc[i, 'chap'] = chapnum
del raw['head']
del raw['mid']
del raw['len']
tmpchap = raw[raw['chap'] == 1].copy()
tmpchap.reset_index(drop=True, inplace=True)
tmpchap['paraidx'] = tmpchap.index
txt = tmpchap.txt[1:]
texts = txt.sum()
# print(texts)
texts = psg.lcut(texts)# 附加词性的分词结果
text = []
word_list = []
stopwords = [line.strip() for line in open('E:\中文文本数据挖掘\停用词.txt', encoding='UTF-8').readlines()]
for item in texts:
    if item.word in stopwords:
        continue
    text.append(item)
    word_list.append(item.word)

# print(text)
name = []
place = []
for item in text:
    # print(item.word, item.flag)
    if len(item.word)>1:
        if item.flag.startswith('nr',):
            name.append(item.word)
        elif item.flag.startswith('ns',):
            place.append(item.word)
        else:
            pass
    else:
        pass
# print(name)
# print(place)

nameDict = {}
for item in name:
    if item in nameDict.keys():
        nameDict[item] += 1
    else:
        nameDict[item] = 1
# print(nameDict)
placeDict = {}
for item in place:
    if item in placeDict.keys():
        placeDict[item] += 1
    else:
        placeDict[item] = 1
# print(placeDict)
#将所有的人名按照蓝色系，地名按照红色系进行词云绘制
myfont = '‪C:\Windows\Fonts\STSONG.TTF'
class GroupedColorFunc(object):
    def __init__(self, color_to_words, default_color):
        self.color_func_to_words = [
            (get_single_color_func(color), set(words))
            for (color, words) in color_to_words.items()]
        self.default_color_func = get_single_color_func(default_color)

    def get_color_func(self, word):
        """Returns a single_color_func associated with the word"""
        try:
            color_func = next(
                color_func for (color_func, words) in self.color_func_to_words
                if word in words)
        except StopIteration:
            color_func = self.default_color_func

        return color_func

    def __call__(self, word, **kwargs):
        return self.get_color_func(word)(word, **kwargs)

######
# 指定分组色系
color_to_words = {
    'red': place,
    'blue': name
}
default_color = 'grey' # 指定其他词条的颜色
grouped_color_func = GroupedColorFunc(color_to_words, default_color)
cloudobj = wordcloud.WordCloud(font_path = myfont,
    width = 1200, height = 800,
    mode = "RGBA", background_color = None).generate(' '.join(word_list))
cloudobj.recolor(color_func=grouped_color_func)
plt.imshow(cloudobj)
plt.axis("off")
plt.show()
cloudobj.to_image()
cloudobj.to_file("人名地名词云图绘制.png")

#制作两个纯色图片，分别为绿色和蓝色
#绿色
word_cloud = wordcloud.WordCloud(
    font_path=myfont,
    background_color='white',
    mask=imread('1.jpg'),
    mode="RGBA",
    colormap="Greens",
    contour_color='steelblue',
).fit_words(nameDict)
plt.imshow(word_cloud)
plt.axis("off")
plt.show()
word_cloud.to_image()
word_cloud.to_file("绿色纯色系绘图.png")
#蓝色
word_cloud = wordcloud.WordCloud(
    font_path=myfont,
    background_color='white',
    mask=imread('1.jpg'),
    mode="RGBA",
    colormap="Blues",
    contour_color='steelblue',
).fit_words(nameDict)
plt.imshow(word_cloud)
plt.axis("off")
plt.show()
word_cloud.to_image()
word_cloud.to_file("蓝色纯色系绘图.png")
